// 并查集
#pragma once

#include <vector>
#include <map>
#include <utility>
#include <iostream>

namespace UnionFind
{
    template <typename T>
    class UnionFind
    {
    public:
        UnionFind(std::vector<T> vertex)
        :_vertex(vertex)
        {
            int n = vertex.size();
            _index = std::vector<int>(n, -1);
            for (int i = 0; i < n; i++)
            {
                _hash[vertex[i]] = i;
            }
        }

        ~UnionFind()
        {
        }

        // 合并e1和e2
        bool UnionSet(T element1, T element2)
        {
            // 找到两个元素所在树的根，根的节点小的
            int root1 = FindRoot(element1);
            int root2 = FindRoot(element2);

            if (root1 == root2)
            {
                return false;
            }
            if (abs(_index[root1]) < abs(_index[root2]))
                std::swap(root1, root2);
            _index[root1] += _index[root2];
            _index[root2] = root1;
            return true;
        }

        int FindRoot(T e)
        {
            // 优化路径压缩
            // 如果是根就直接返回，如果遍历到根的路径，并把路径上的所有结点都连接到根上
            if (_index[_hash[e]] < 0)
                return _hash[e];
            int idx = _hash[e];
            while (_index[idx] >= 0)
            {
                idx = _index[idx];
            }
            int i = _hash[e];
            while (i != idx)
            {
                int tmp = _index[i];
                _index[i] = idx;
                i = tmp;
            }
            return idx;
        }

        bool InSameSet(T element1, T element2)
        {
            return FindRoot(element1) == FindRoot(element2);
        }

        int TreeNum()
        {
            int count = 0;
            for (auto e : _index)
            {
                if (e < 0)
                    count++;
            }
            return count;
        }

        //值找下标
        int FindIndex(T e)
        {
            auto t = _hash.find(e);
            if (t != _hash.end())
            {
                return t->second;
            }
            else
            {
                // throw std::invalid_argument("顶点不存在");
				return -1;
            }
        }

        //下标找值
        T FindValue(int index)
        {
            //通过下标找值
            return _vertex[index];
        }

        int Size()
        {
            return _vertex.size();
        }

    private:
        std::vector<T> _vertex; //顶点集合
        std::vector<int> _index; // 下标集合
        std::map<T, int> _hash;  // 顶点与下标的映射
    };
}

typedef UnionFind::UnionFind<std::string> uf_str;

void TestUnionFind()
{
    std::vector<std::string> name{"张三", "李四", "王五", "李阳", "阳sir"};
    uf_str uf(name);
    uf.UnionSet("李阳", "阳sir");
    std::cout << uf.InSameSet("李阳", "阳sir") << std::endl;
}
